{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/amax/anaconda3/envs/mopo/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:516: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/amax/anaconda3/envs/mopo/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:517: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/amax/anaconda3/envs/mopo/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:518: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/amax/anaconda3/envs/mopo/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:519: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/amax/anaconda3/envs/mopo/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:520: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/amax/anaconda3/envs/mopo/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n",
      "Warning: Flow failed to import. Set the environment variable D4RL_SUPPRESS_IMPORT_ERROR=1 to suppress this message.\n",
      "No module named 'flow'\n",
      "Warning: CARLA failed to import. Set the environment variable D4RL_SUPPRESS_IMPORT_ERROR=1 to suppress this message.\n",
      "No module named 'carla'\n"
     ]
    }
   ],
   "source": [
    "from mopo.models.cross_policy_evaluate import load_model, get_mask, cross_policy_evaluate\n",
    "import gym\n",
    "import d4rl\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import mopo.mask\n",
    "import os\n",
    "import shutil\n",
    "import argparse\n",
    "from RLA.easy_log.tester import tester\n",
    "import json\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_name = \"ray_mopo\"\n",
    "origin_name = \"hopper\"\n",
    "target_name = \"hopper_baselines\"\n",
    "\n",
    "dataset_name = \"hopper_medium_replay_200e3\"\n",
    "seeds = [\"1762\", \"5754\", \"5955\", \"995\"]\n",
    "\n",
    "\n",
    "os.makedirs(os.path.join(root_name, target_name, dataset_name), exist_ok=True)\n",
    "all_names = os.listdir(os.path.join(root_name, origin_name, dataset_name))\n",
    "for seed in seeds:\n",
    "    l = list(filter(lambda x: x.startswith(\"seed:\" + seed + \"_\"), all_names))\n",
    "    if len(l) == 1:\n",
    "        seed_full_name = l[0]\n",
    "    else:\n",
    "        print(\"wrong\", l)\n",
    "        continue\n",
    "    data_name = seed_full_name[len(\"seed:\" + seed) + 1:-8]\n",
    "    experiment_name = \"experiment_state-\" + data_name + \".json\"\n",
    "    shutil.move(os.path.join(root_name, origin_name, dataset_name, seed_full_name), \n",
    "               os.path.join(root_name, target_name, dataset_name, seed_full_name))\n",
    "    shutil.move(os.path.join(root_name, origin_name, dataset_name, experiment_name), \n",
    "           os.path.join(root_name, target_name, dataset_name, experiment_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6367264607445862"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from d4rl import ope\n",
    "\n",
    "ope.normalize(\"hopper-medium-replay\", 2052)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_exp_dir=\"hopper\"\n",
    "dataset=\"hopper_medium_replay_200e3\" \n",
    "query = dict(algorithm_params=dict(\n",
    "                 network_structure=\"shared\",\n",
    "                 mask=\"dense\",\n",
    "                 kwargs=dict(\n",
    "                     hidden_dim=100,\n",
    "                     )\n",
    "                 )\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hopper-medium-replay-v0\n"
     ]
    }
   ],
   "source": [
    "def is_dict_consist(param, query):\n",
    "    if isinstance(query, dict):\n",
    "        consist = True\n",
    "        for key in query:\n",
    "            consist = consist and is_dict_consist(param[key], query[key])\n",
    "        return consist\n",
    "    else:\n",
    "        return param == query\n",
    "    \n",
    "def dataset_name_to_mask_name(dataset_name):\n",
    "    l = dataset_name.split(\"_\")[:-1] + [\"v0\"]\n",
    "    return \"-\".join(l)\n",
    "print(dataset_name_to_mask_name(dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cross_policy_evaluate_from_mopo(base_exp_dir, dataset, query):\n",
    "    all_path = os.listdir(os.path.join(\"ray_mopo\", base_exp_dir, dataset))\n",
    "    all_log_path = list(filter(lambda x: x.startswith(\"seed\"), all_path))\n",
    "    \n",
    "    train_env_name = dataset_name_to_mask_name(dataset)\n",
    "    test_env_name = train_env_name\n",
    "    train_losses = []\n",
    "    test_losses = []\n",
    "    \n",
    "    mask = get_mask(dataset.split(\"_\")[0], query[\"algorithm_params\"][\"mask\"])\n",
    "    \n",
    "    for log_path in all_log_path:\n",
    "        param_path = os.path.join(\"ray_mopo\", base_exp_dir, dataset, log_path, \"params.json\")\n",
    "        with open(param_path) as f:\n",
    "            param = json.loads(f.read())\n",
    "            if is_dict_consist(param, query):\n",
    "                model_save_dir = os.path.join(\"ray_mopo\", base_exp_dir, dataset, log_path, \"models\")\n",
    "                train_loss, test_loss = cross_policy_evaluate(train_env_name, test_env_name, \n",
    "                                          query[\"algorithm_params\"][\"network_structure\"],\n",
    "                                          mask,\n",
    "                                          model_save_dir, \n",
    "                                          hidden_dim=query[\"algorithm_params\"][\"kwargs\"][\"hidden_dim\"]) \n",
    "                train_losses.append(train_loss)\n",
    "                test_losses.append(test_loss)\n",
    "    return np.array(train_losses), np.array(test_losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/amax/anaconda3/envs/mopo/lib/python3.6/site-packages/gym/logger.py:30: UserWarning: \u001b[33mWARN: Box bound precision lowered by casting to float32\u001b[0m\n",
      "  warnings.warn(colorize('%s: %s'%('WARN', msg % args), 'yellow'))\n",
      "load datafile: 100%|██████████| 5/5 [00:00<00:00, 34.48it/s]\n",
      "load datafile: 100%|██████████| 5/5 [00:00<00:00, 47.16it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ BNN ] Name BNN_0 | Observation dim 11 | Action dim: 3 | Hidden dim: 100\n",
      "Specified load dir ray_mopo/hopper/hopper_medium_replay_200e3/seed:1762_2022-01-18_00-33-48nrvnpg8g/models\n",
      "[ BNN ] Initializing model: BNN_0 | 7 networks | 5 elites\n",
      "Created an ensemble of 7 neural networks with variance predictions | Elites: 5\n",
      "Added layer with input dim 14 , output dim 100\n",
      "Added layer with input dim 100 , output dim 100\n",
      "Added layer with input dim 100 , output dim 100\n",
      "Added layer with input dim 100 , output dim 12\n",
      "Added layer with input dim 100 , output dim 12\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "get loss:   0%|          | 0/10 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ BNN ] Model: <mopo.models.shared_causal_bnn.SharedCausalBNN object at 0x7f6ee58c99b0>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "get loss: 100%|██████████| 10/10 [00:02<00:00,  5.03it/s]\n",
      "get loss: 100%|██████████| 10/10 [00:02<00:00,  4.98it/s]\n",
      "/home/amax/anaconda3/envs/mopo/lib/python3.6/site-packages/gym/logger.py:30: UserWarning: \u001b[33mWARN: Box bound precision lowered by casting to float32\u001b[0m\n",
      "  warnings.warn(colorize('%s: %s'%('WARN', msg % args), 'yellow'))\n",
      "load datafile: 100%|██████████| 5/5 [00:00<00:00, 34.67it/s]\n",
      "load datafile: 100%|██████████| 5/5 [00:00<00:00, 47.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ BNN ] Name BNN_0 | Observation dim 11 | Action dim: 3 | Hidden dim: 100\n",
      "Specified load dir ray_mopo/hopper/hopper_medium_replay_200e3/seed:5754_2022-01-18_00-33-33z59k4z1k/models\n",
      "[ BNN ] Initializing model: BNN_0 | 7 networks | 5 elites\n",
      "Created an ensemble of 7 neural networks with variance predictions | Elites: 5\n",
      "Added layer with input dim 14 , output dim 100\n",
      "Added layer with input dim 100 , output dim 100\n",
      "Added layer with input dim 100 , output dim 100\n",
      "Added layer with input dim 100 , output dim 12\n",
      "Added layer with input dim 100 , output dim 12\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "get loss:   0%|          | 0/10 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ BNN ] Model: <mopo.models.shared_causal_bnn.SharedCausalBNN object at 0x7f6eda0f9780>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "get loss: 100%|██████████| 10/10 [00:02<00:00,  5.07it/s]\n",
      "get loss: 100%|██████████| 10/10 [00:02<00:00,  5.06it/s]\n",
      "/home/amax/anaconda3/envs/mopo/lib/python3.6/site-packages/gym/logger.py:30: UserWarning: \u001b[33mWARN: Box bound precision lowered by casting to float32\u001b[0m\n",
      "  warnings.warn(colorize('%s: %s'%('WARN', msg % args), 'yellow'))\n",
      "load datafile: 100%|██████████| 5/5 [00:00<00:00, 46.00it/s]\n",
      "load datafile: 100%|██████████| 5/5 [00:00<00:00, 47.86it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ BNN ] Name BNN_0 | Observation dim 11 | Action dim: 3 | Hidden dim: 100\n",
      "Specified load dir ray_mopo/hopper/hopper_medium_replay_200e3/seed:5955_2022-01-18_00-32-563aewe8hc/models\n",
      "[ BNN ] Initializing model: BNN_0 | 7 networks | 5 elites\n",
      "Created an ensemble of 7 neural networks with variance predictions | Elites: 5\n",
      "Added layer with input dim 14 , output dim 100\n",
      "Added layer with input dim 100 , output dim 100\n",
      "Added layer with input dim 100 , output dim 100\n",
      "Added layer with input dim 100 , output dim 12\n",
      "Added layer with input dim 100 , output dim 12\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "get loss:   0%|          | 0/10 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ BNN ] Model: <mopo.models.shared_causal_bnn.SharedCausalBNN object at 0x7f6e78a44630>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "get loss: 100%|██████████| 10/10 [00:02<00:00,  4.55it/s]\n",
      "get loss: 100%|██████████| 10/10 [00:02<00:00,  4.68it/s]\n",
      "/home/amax/anaconda3/envs/mopo/lib/python3.6/site-packages/gym/logger.py:30: UserWarning: \u001b[33mWARN: Box bound precision lowered by casting to float32\u001b[0m\n",
      "  warnings.warn(colorize('%s: %s'%('WARN', msg % args), 'yellow'))\n",
      "load datafile: 100%|██████████| 5/5 [00:00<00:00, 34.76it/s]\n",
      "load datafile: 100%|██████████| 5/5 [00:00<00:00, 47.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ BNN ] Name BNN_0 | Observation dim 11 | Action dim: 3 | Hidden dim: 100\n",
      "Specified load dir ray_mopo/hopper/hopper_medium_replay_200e3/seed:995_2022-01-18_00-33-01h8p87vco/models\n",
      "[ BNN ] Initializing model: BNN_0 | 7 networks | 5 elites\n",
      "Created an ensemble of 7 neural networks with variance predictions | Elites: 5\n",
      "Added layer with input dim 14 , output dim 100\n",
      "Added layer with input dim 100 , output dim 100\n",
      "Added layer with input dim 100 , output dim 100\n",
      "Added layer with input dim 100 , output dim 12\n",
      "Added layer with input dim 100 , output dim 12\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "get loss:   0%|          | 0/10 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ BNN ] Model: <mopo.models.shared_causal_bnn.SharedCausalBNN object at 0x7f6ee58abc18>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "get loss: 100%|██████████| 10/10 [00:02<00:00,  5.10it/s]\n",
      "get loss: 100%|██████████| 10/10 [00:02<00:00,  5.14it/s]\n"
     ]
    }
   ],
   "source": [
    "train_losses, test_losses = cross_policy_evaluate_from_mopo(base_exp_dir, dataset, query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.23299e-04, 1.49065e-03, 8.96284e-05, 8.73359e-03, 5.70373e-05,\n",
       "       2.61320e-03, 3.37344e-02, 2.51611e-02, 2.18192e-02, 5.46402e-02,\n",
       "       2.04212e-02, 1.41445e-01])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_losses.mean(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.11301e-04, 6.56411e-04, 5.68515e-05, 4.56854e-03, 4.24999e-05,\n",
       "       1.70904e-03, 2.50216e-02, 1.59172e-02, 1.16844e-02, 2.75748e-02,\n",
       "       9.76517e-03, 7.93953e-02])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cross_policy_evaluate_result[\"none\"]['medium-replay'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([9.03134e-04, 1.54911e-03, 1.05910e-04, 7.26656e-03, 4.79932e-05,\n",
       "       2.65328e-03, 3.62899e-02, 4.25683e-02, 3.10165e-02, 6.02638e-02,\n",
       "       2.73303e-02, 2.10813e-01])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cross_policy_evaluate_result[\"dense\"]['medium-replay'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.18129e-03, 1.33646e-03, 8.63915e-05, 6.58226e-03, 6.09931e-05,\n",
       "       2.66039e-03, 3.54992e-02, 4.23978e-02, 2.67683e-02, 5.17967e-02,\n",
       "       2.02406e-02, 1.03869e+00])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cross_policy_evaluate_result[\"binary_04_medium_replay_causal\"]['medium-replay'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([5.01848e-03, 3.49361e-03, 1.33565e-04, 9.07713e-03, 1.92046e-04,\n",
       "       2.89930e-03, 3.99379e-02, 4.15243e-02, 5.91794e-02, 8.78187e-02,\n",
       "       2.23160e-02, 4.17495e-01])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cross_policy_evaluate_result[\"binary_03_medium_replay_causal\"]['medium-replay'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3.53922e-02, 5.88904e-02, 1.18510e-04, 8.81222e-03, 1.87134e-04,\n",
       "       3.07766e-03, 6.49895e+00, 1.16227e+01, 4.66007e+00, 1.34766e-01,\n",
       "       1.66958e-01, 1.05643e+00])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cross_policy_evaluate_result[\"binary_01_medium_replay_causal\"]['medium-replay'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.69381e-02, 8.26936e-03, 1.07804e-04, 8.78500e-03, 6.38726e-05,\n",
       "       2.86847e-03, 2.73332e-01, 8.90458e-01, 6.92159e-02, 1.04552e-01,\n",
       "       6.22180e-02, 1.01987e+00])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cross_policy_evaluate_result[\"binary_005_medium_replay_causal\"]['medium-replay'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
